/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "parcsr_interpolation.h"
#include "parcsr_matmul.h"

namespace YHAMG
{

#define FINE -1

#define ABS(x) ((x) > 0 ? (x) : -(x))
#define CMPABS(a, b) (ABS(a) > ABS(b))
#define _(x) ABS(x) // ((x) > 0 ^ diagonal > 0 ? (x) : 0.0)

static inline void Swap(int& a, int& b)
{
	int temp = a;
	a = b;
	b = temp;
}

static inline void Swap(double& a, double& b)
{
	double temp = a;
	a = b;
	b = temp;
}

static void SortAbs(int* a, double* x, int left, int right)
{
	if (left >= right) return;

	Swap(a[left], a[(left + right) / 2]);
	Swap(x[left], x[(left + right) / 2]);

	int last = left;
	for (int i = left + 1; i <= right; ++i)
	{
		if (CMPABS(x[i], x[left]))
		{
			Swap(a[++last], a[i]);
			Swap(x[last], x[i]);
		}
	}

	Swap(a[left], a[last]);
	Swap(x[left], x[last]);

	SortAbs(a, x, left, last - 1);
	SortAbs(a, x, last + 1, right);
}

#define MPI_TAG 450

void ParCSR_Interpolation(const ParCSR_Matrix& A, const bool* LocLoc_s, const bool* LocExt_s, int* CF_map, ParCSR_Matrix& P, int iterations, int min_elements, int max_elements, double truncation_factor)
{
	MPI_Comm comm = A.comm;

	int n = A.LocLoc.size[0];
	int Ext_size = A.LocExt.size[1];

	int* A_LocLoc_rowptr = A.LocLoc.rowptr;
	int* A_LocLoc_colind = A.LocLoc.colind;
	double* A_LocLoc_values = A.LocLoc.values;
	int* A_LocExt_rowptr = A.LocExt.rowptr;
	int* A_LocExt_colind = A.LocExt.colind;
	double* A_LocExt_values = A.LocExt.values;

	int num_neighbors = A.num_neighbors;
	int* neighbor_ranks = A.neighbor_ranks;

	int* Ext_ptr = A.Ext_ptr;
	int* Ext_map = A.Ext_map;

	int* recv_ptr = A.Ext_ptr;
	int* send_ptr = A.send_ptr;
	int* send_list = A.send_list;
	int* send_buffer = new int[send_ptr[num_neighbors]];

	int num_coarse = 0;
	for (int i = 0; i < n; ++i)
		if (CF_map[i] != FINE)
			CF_map[i] = num_coarse++;

	int* Ext_CF_map = new int[Ext_size];

	MPI_Request* recv_requests = new MPI_Request[num_neighbors];
	MPI_Request* send_requests = new MPI_Request[num_neighbors];
	MPI_Status status;

	for (int r = 0; r < num_neighbors; ++r)
		if (recv_ptr[r + 1] > recv_ptr[r])
			MPI_Irecv(Ext_CF_map + recv_ptr[r], recv_ptr[r + 1] - recv_ptr[r],
				MPI_INT, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
				send_buffer[i] = CF_map[send_list[i]];
			MPI_Isend(send_buffer + send_ptr[r], send_ptr[r + 1] - send_ptr[r], MPI_INT, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);
		}
	}

	for (int r = 0; r < num_neighbors; ++r)
		if (recv_ptr[r + 1] > recv_ptr[r])
			MPI_Wait(recv_requests + r, &status);

	for (int r = 0; r < num_neighbors; ++r)
		if (send_ptr[r + 1] > send_ptr[r])
			MPI_Wait(send_requests + r, &status);

	delete[] recv_requests;
	delete[] send_requests;
	delete[] send_buffer;
	
	int* Q_Ext_map = new int[Ext_size];
	int* Q_Ext_ptr = new int[num_neighbors + 1];

	Q_Ext_ptr[0] = 0;
	for (int r = 0, k = 0; r < num_neighbors; Q_Ext_ptr[++r] = k)
	{
		for (int i = recv_ptr[r]; i < recv_ptr[r + 1]; ++i)
		{
			if (Ext_CF_map[i] != FINE)
			{
				Q_Ext_map[k] = Ext_CF_map[i];
				Ext_CF_map[i] = k++;
			}
		}
	}
	
	bool* marks = new bool[n];
	double* scalars = new double[n];

	int* Q_LocLoc_rowptr = new int[n + 1];
	int* Q_LocExt_rowptr = new int[n + 1];
	int* R_LocLoc_rowptr = new int[n + 1];
	int* R_LocExt_rowptr = new int[n + 1];

	Q_LocLoc_rowptr[0] = 0;
	Q_LocExt_rowptr[0] = 0;
	R_LocLoc_rowptr[0] = 0;
	R_LocExt_rowptr[0] = 0;

#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
	{
		if (CF_map[i] == FINE)
		{
			int cnt_loc_cs = 0;
			int cnt_loc_n = 0;
			int cnt_ext_cs = 0;
			int cnt_ext_n = 0;

			for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
			{
				int jcol = A_LocLoc_colind[j];
				if (jcol == i) continue;
				++cnt_loc_n;
				if (LocLoc_s[j] && CF_map[jcol] != FINE) ++cnt_loc_cs;
			}

			for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
			{
				int jcol = A_LocExt_colind[j];
				++cnt_ext_n;
				if (LocExt_s[j] && Ext_CF_map[jcol] != FINE) ++cnt_ext_cs;
			}

			marks[i] = cnt_loc_cs + cnt_ext_cs >= min_elements && (cnt_loc_n > cnt_loc_cs || cnt_ext_n > cnt_ext_cs);
			Q_LocLoc_rowptr[i + 1] = cnt_loc_cs;
			Q_LocExt_rowptr[i + 1] = cnt_ext_cs;
			R_LocLoc_rowptr[i + 1] = marks[i] ? 1 : cnt_loc_n;
			R_LocExt_rowptr[i + 1] = marks[i] ? 0 : cnt_ext_n;
		}
		else
		{
			marks[i] = 1;
			Q_LocLoc_rowptr[i + 1] = 1;
			Q_LocExt_rowptr[i + 1] = 0;
			R_LocLoc_rowptr[i + 1] = 1;
			R_LocExt_rowptr[i + 1] = 0;
		}
	}

	for (int i = 0; i < n; ++i)
		Q_LocLoc_rowptr[i + 1] += Q_LocLoc_rowptr[i];
	for (int i = 0; i < n; ++i)
		Q_LocExt_rowptr[i + 1] += Q_LocExt_rowptr[i];
	for (int i = 0; i < n; ++i)
		R_LocLoc_rowptr[i + 1] += R_LocLoc_rowptr[i];
	for (int i = 0; i < n; ++i)
		R_LocExt_rowptr[i + 1] += R_LocExt_rowptr[i];

	int* Q_LocLoc_colind = new int[Q_LocLoc_rowptr[n]];
	double* Q_LocLoc_values = new double[Q_LocLoc_rowptr[n]];
	int* Q_LocExt_colind = new int[Q_LocExt_rowptr[n]];
	double* Q_LocExt_values = new double[Q_LocExt_rowptr[n]];
	int* R_LocLoc_colind = new int[R_LocLoc_rowptr[n]];
	double* R_LocLoc_values = new double[R_LocLoc_rowptr[n]];
	int* R_LocExt_colind = new int[R_LocExt_rowptr[n]];
	double* R_LocExt_values = new double[R_LocExt_rowptr[n]];

#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
	{
		if (CF_map[i] == FINE)
		{
			if (marks[i])
			{
				double diagonal = 0.0;
				double sum_n = 0.0;

				for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					if (A_LocLoc_colind[j] == i)
						diagonal = A_LocLoc_values[j];
					else
						sum_n += A_LocLoc_values[j];
				}

				for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
					sum_n += A_LocExt_values[j];

				double alpha = diagonal != 0 ? -sum_n / diagonal : 0.0;
				
				double sum_q = 0.0;

				for (int j = A_LocLoc_rowptr[i], k = Q_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocLoc_colind[j];

					if (LocLoc_s[j] && CF_map[jcol] != FINE)
					{
						double _jval = _(A_LocLoc_values[j]);
						Q_LocLoc_colind[k] = CF_map[jcol];
						Q_LocLoc_values[k++] = _jval;
						sum_q += _jval;
					}
				}

				for (int j = A_LocExt_rowptr[i], k = Q_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocExt_colind[j];

					if (LocExt_s[j] && Ext_CF_map[jcol] != FINE)
					{
						double _jval = _(A_LocExt_values[j]);
						Q_LocExt_colind[k] = Ext_CF_map[jcol];
						Q_LocExt_values[k++] = _jval;
						sum_q += _jval;
					}
				}

				if (sum_q != 0)
				{
					double scale = alpha / sum_q;
					for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j)
						Q_LocLoc_values[j] *= scale;
					for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j)
						Q_LocExt_values[j] *= scale;
				}
				else
				{
					for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j)
						Q_LocLoc_values[j] = 0.0;
					for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j)
						Q_LocExt_values[j] = 0.0;
				}

				R_LocLoc_colind[R_LocLoc_rowptr[i]] = i;
				R_LocLoc_values[R_LocLoc_rowptr[i]] = 1.0;

				scalars[i] = alpha;
			}
			else
			{
				double diagonal = 0.0;
				double sum_n = 0.0;

				for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					if (A_LocLoc_colind[j] == i)
						diagonal = A_LocLoc_values[j];
					else
						sum_n += A_LocLoc_values[j];
				}

				for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
					sum_n += A_LocExt_values[j];

				double alpha = diagonal != 0 ? -sum_n / diagonal : 0.0;

				double sum_r = 0.0;
				double sum_q = 0.0;

				for (int j = A_LocLoc_rowptr[i], k = Q_LocLoc_rowptr[i], r = R_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocLoc_colind[j];

					if (jcol == i) continue;

					double _jval = _(A_LocLoc_values[j]);
					
					R_LocLoc_colind[r] = jcol;
					R_LocLoc_values[r++] = _jval;
					sum_r += _jval;

					if (LocLoc_s[j] && CF_map[jcol] != FINE)
					{
						Q_LocLoc_colind[k] = CF_map[jcol];
						Q_LocLoc_values[k++] = _jval;
						sum_q += _jval;
					}
				}

				for (int j = A_LocExt_rowptr[i], k = Q_LocExt_rowptr[i], r = R_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
				{
					int jcol = A_LocExt_colind[j];
					double _jval = _(A_LocExt_values[j]);

					R_LocExt_colind[r] = jcol;
					R_LocExt_values[r++] = _jval;
					sum_r += _jval;

					if (LocExt_s[j] && Ext_CF_map[jcol] != FINE)
					{
						Q_LocExt_colind[k] = Ext_CF_map[jcol];
						Q_LocExt_values[k++] = _jval;
						sum_q += _jval;
					}
				}

				if (sum_q != 0)
				{
					double scale = alpha / sum_q;
					for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j)
						Q_LocLoc_values[j] *= scale;
					for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j)
						Q_LocExt_values[j] *= scale;
				}
				else
				{
					for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j)
						Q_LocLoc_values[j] = 0.0;
					for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j)
						Q_LocExt_values[j] = 0.0;
				}

				if (sum_r != 0)
				{
					double scale = alpha / sum_r;
					for (int j = R_LocLoc_rowptr[i]; j < R_LocLoc_rowptr[i + 1]; ++j)
						R_LocLoc_values[j] *= scale;
					for (int j = R_LocExt_rowptr[i]; j < R_LocExt_rowptr[i + 1]; ++j)
						R_LocExt_values[j] *= scale;
				}
				else
				{
					for (int j = R_LocLoc_rowptr[i]; j < R_LocLoc_rowptr[i + 1]; ++j)
						R_LocLoc_values[j] = 0.0;
					for (int j = R_LocExt_rowptr[i]; j < R_LocExt_rowptr[i + 1]; ++j)
						R_LocExt_values[j] = 0.0;
				}

				scalars[i] = alpha;
			}
		}
		else
		{
			Q_LocLoc_colind[Q_LocLoc_rowptr[i]] = CF_map[i];
			Q_LocLoc_values[Q_LocLoc_rowptr[i]] = 1.0;
			R_LocLoc_colind[R_LocLoc_rowptr[i]] = i;
			R_LocLoc_values[R_LocLoc_rowptr[i]] = 1.0;
		}
	}

	int* Q_neighbors = new int[num_neighbors];
	for (int r = 0; r < num_neighbors; ++r)
		Q_neighbors[r] = neighbor_ranks[r];

	int* R_neighbors = new int[num_neighbors];
	for (int r = 0; r < num_neighbors; ++r)
		R_neighbors[r] = neighbor_ranks[r];

	int* v = new int[Ext_size];
	for (int i = 0; i < Ext_size; ++i)
		v[i] = -1;

	for (int j = R_LocExt_rowptr[0]; j < R_LocExt_rowptr[n]; ++j)
		v[R_LocExt_colind[j]] = 1;

	int* R_Ext_map = new int[Ext_size];
	int* R_Ext_ptr = new int[num_neighbors + 1];

	R_Ext_ptr[0] = 0;
	for (int r = 0, k = 0; r < num_neighbors; R_Ext_ptr[++r] = k)
	{
		for (int i = Ext_ptr[r]; i < Ext_ptr[r + 1]; ++i)
		{
			if (v[i] == 1)
			{
				v[i] = k;
				R_Ext_map[k++] = Ext_map[i];
			}
		}
	}

	for (int j = R_LocExt_rowptr[0]; j < R_LocExt_rowptr[n]; ++j)
		R_LocExt_colind[j] = v[R_LocExt_colind[j]];

	delete[] v;

	ParCSR_Matrix R, Q;

	R.comm = comm;
	R.LocLoc.size[0] = n;
	R.LocLoc.size[1] = n;
	R.LocLoc.rowptr = R_LocLoc_rowptr;
	R.LocLoc.colind = R_LocLoc_colind;
	R.LocLoc.values = R_LocLoc_values;

	R.LocExt.size[0] = n;
	R.LocExt.size[1] = R_Ext_ptr[num_neighbors];
	R.LocExt.rowptr = R_LocExt_rowptr;
	R.LocExt.colind = R_LocExt_colind;
	R.LocExt.values = R_LocExt_values;

	R.num_neighbors = num_neighbors;
	R.neighbor_ranks = R_neighbors;
	R.Ext_ptr = R_Ext_ptr;
	R.Ext_map = R_Ext_map;

	Q.comm = comm;
	Q.LocLoc.size[0] = n;
	Q.LocLoc.size[1] = num_coarse;
	Q.LocLoc.rowptr = Q_LocLoc_rowptr;
	Q.LocLoc.colind = Q_LocLoc_colind;
	Q.LocLoc.values = Q_LocLoc_values;

	Q.LocExt.size[0] = n;
	Q.LocExt.size[1] = Q_Ext_ptr[num_neighbors];
	Q.LocExt.rowptr = Q_LocExt_rowptr;
	Q.LocExt.colind = Q_LocExt_colind;
	Q.LocExt.values = Q_LocExt_values;

	Q.num_neighbors = num_neighbors;
	Q.neighbor_ranks = Q_neighbors;
	Q.Ext_ptr = Q_Ext_ptr;
	Q.Ext_map = Q_Ext_map;

	R.SetupHalo();

	for (int iter = 0; iter < iterations; ++iter)
	{
		if (iter > 1)
		{
			Q_LocLoc_rowptr = Q.LocLoc.rowptr;
			Q_LocLoc_colind = Q.LocLoc.colind;
			Q_LocLoc_values = Q.LocLoc.values;

			Q_LocExt_rowptr = Q.LocExt.rowptr;
			Q_LocExt_colind = Q.LocExt.colind;
			Q_LocExt_values = Q.LocExt.values;

	#ifdef USE_OPENMP
	#pragma omp for schedule(guided)
	#endif
			for (int i = 0; i < n; ++i)
			{
				if (!marks[i])
				{
					double sum = 0.0;

					for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j)
						sum += Q_LocLoc_values[j];
					for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j)
						sum += Q_LocExt_values[j];

					if (sum != 0)
					{
						double scale = scalars[i] / sum;
						for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j)
							Q_LocLoc_values[j] *= scale;
						for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j)
							Q_LocExt_values[j] *= scale;
					}
					else
					{
						for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j)
							Q_LocLoc_values[j] = 0.0;
						for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j)
							Q_LocExt_values[j] = 0.0;
					}
				}
			}
			
			for (int i = 0; i < n; ++i)
				if (!marks[i])
					marks[i] = Q_LocLoc_rowptr[i + 1] - Q_LocLoc_rowptr[i]
					+ Q_LocExt_rowptr[i + 1] - Q_LocExt_rowptr[i] >= min_elements;

			R_LocLoc_rowptr = R.LocLoc.rowptr;
			R_LocLoc_colind = R.LocLoc.colind;
			R_LocLoc_values = R.LocLoc.values;

			R_LocExt_rowptr = R.LocExt.rowptr;
			R_LocExt_colind = R.LocExt.colind;
			R_LocExt_values = R.LocExt.values;

			int* _R_LocLoc_rowptr = new int[n + 1];
			int* _R_LocExt_rowptr = new int[n + 1];

			_R_LocLoc_rowptr[0] = 0;
			_R_LocExt_rowptr[0] = 0;
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				if (!marks[i])
				{
					_R_LocLoc_rowptr[i + 1] = R_LocLoc_rowptr[i + 1] - R_LocLoc_rowptr[i];
					_R_LocExt_rowptr[i + 1] = R_LocExt_rowptr[i + 1] - R_LocExt_rowptr[i];
				}
				else
				{
					_R_LocLoc_rowptr[i + 1] = 1;
					_R_LocExt_rowptr[i + 1] = 0;
				}
			}

			for (int i = 0; i < n; ++i)
				_R_LocLoc_rowptr[i + 1] += _R_LocLoc_rowptr[i];
			for (int i = 0; i < n; ++i)
				_R_LocExt_rowptr[i + 1] += _R_LocExt_rowptr[i];

			int* _R_LocLoc_colind = new int[_R_LocLoc_rowptr[n]];
			double* _R_LocLoc_values = new double[_R_LocLoc_rowptr[n]];
			int* _R_LocExt_colind = new int[_R_LocExt_rowptr[n]];
			double* _R_LocExt_values = new double[_R_LocExt_rowptr[n]];
#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
			for (int i = 0; i < n; ++i)
			{
				if (!marks[i])
				{
					for (int j = R_LocLoc_rowptr[i], k = _R_LocLoc_rowptr[i]; j < R_LocLoc_rowptr[i + 1]; ++j, ++k)
					{
						_R_LocLoc_colind[k] = R_LocLoc_colind[j];
						_R_LocLoc_values[k] = R_LocLoc_values[j];
					}
					for (int j = R_LocExt_rowptr[i], k = _R_LocExt_rowptr[i]; j < R_LocExt_rowptr[i + 1]; ++j, ++k)
					{
						_R_LocExt_colind[k] = R_LocExt_colind[j];
						_R_LocExt_values[k] = R_LocExt_values[j];
					}
				}
				else
				{
					_R_LocLoc_colind[_R_LocLoc_rowptr[i]] = i;
					_R_LocLoc_values[_R_LocLoc_rowptr[i]] = 1.0;
				}
			}

			int* w = new int[R.LocExt.size[1]];
			for (int i = 0; i < R.LocExt.size[1]; ++i)
				w[i] = -1;

			for (int j = _R_LocExt_rowptr[0]; j < _R_LocExt_rowptr[n]; ++j)
				w[_R_LocExt_colind[j]] = 1;

			for (int r = 0, k = 0, i = 0; r < R.num_neighbors; R_Ext_ptr[++r] = k)
			{
				for ( ; i < R_Ext_ptr[r + 1]; ++i)
				{
					if (w[i] == 1)
					{
						w[i] = k;
						R_Ext_map[k++] = R_Ext_map[i];
					}
				}
			}

			for (int j = _R_LocExt_rowptr[0]; j < _R_LocExt_rowptr[n]; ++j)
				_R_LocExt_colind[j] = w[_R_LocExt_colind[j]];

			delete[] w;

			R.LocLoc.Free();
			R.LocExt.Free();

			R.LocLoc.size[0] = n;
			R.LocLoc.size[1] = n;
			R.LocLoc.rowptr = _R_LocLoc_rowptr;
			R.LocLoc.colind = _R_LocLoc_colind;
			R.LocLoc.values = _R_LocLoc_values;

			R.LocExt.size[0] = n;
			R.LocExt.size[1] = R_Ext_ptr[R.num_neighbors];
			R.LocExt.rowptr = _R_LocExt_rowptr;
			R.LocExt.colind = _R_LocExt_colind;
			R.LocExt.values = _R_LocExt_values;
			
			R.SetupHalo();
		}

		ParCSR_MatMul(R, Q, Q);
	}

	Q_LocLoc_rowptr = Q.LocLoc.rowptr;
	Q_LocLoc_colind = Q.LocLoc.colind;
	Q_LocLoc_values = Q.LocLoc.values;

	Q_LocExt_rowptr = Q.LocExt.rowptr;
	Q_LocExt_colind = Q.LocExt.colind;
	Q_LocExt_values = Q.LocExt.values;

	int Q_num_neighbors = Q.num_neighbors;
	Q_neighbors = Q.neighbor_ranks;
	Q_Ext_ptr = Q.Ext_ptr;
	Q_Ext_map = Q.Ext_map;

	int* P_LocLoc_rowptr = new int[n + 1];
	int* P_LocExt_rowptr = new int[n + 1];

	P_LocLoc_rowptr[0] = 0;
	P_LocExt_rowptr[0] = 0;

#ifdef USE_OPENMP
#pragma omp parallel
#endif
	{
#ifdef USE_OPENMP
#pragma omp for schedule(guided)
#endif
		for (int i = 0; i < n; ++i)
		{
			if (CF_map[i] == FINE)
			{
				int m = Q_LocLoc_rowptr[i + 1] - Q_LocLoc_rowptr[i] + Q_LocExt_rowptr[i + 1] - Q_LocExt_rowptr[i];
				int* colind = new int[m];
				double* values = new double[m];
				
				int r = 0;
				
				for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i + 1]; ++j, ++r)
				{
					colind[r] = Q_LocLoc_colind[j];
					values[r] = Q_LocLoc_values[j];
				}

				for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i + 1]; ++j, ++r)
				{
					colind[r] = Q_LocExt_colind[j] + num_coarse;
					values[r] = Q_LocExt_values[j];
				}

				SortAbs(colind, values, 0, m - 1);

				int cnt_loc = 0;
				int cnt_ext = 0;

				double alpha = truncation_factor * ABS(values[0]);

				for (int j = 0; j < r; ++j)
				{
					if (ABS(values[j]) < alpha || j == max_elements) break;
					
					if (colind[j] < num_coarse)
					{
						Q_LocLoc_colind[Q_LocLoc_rowptr[i] + cnt_loc] = colind[j];
						Q_LocLoc_values[Q_LocLoc_rowptr[i] + cnt_loc++] = values[j];
					}
					else
					{
						Q_LocExt_colind[Q_LocExt_rowptr[i] + cnt_ext] = colind[j] - num_coarse;
						Q_LocExt_values[Q_LocExt_rowptr[i] + cnt_ext++] = values[j];
					}
				}

				delete[] colind;
				delete[] values;

				double sum = 0.0;

				for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i] + cnt_loc; ++j)
					sum += Q_LocLoc_values[j];
				for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i] + cnt_ext; ++j)
					sum += Q_LocExt_values[j];

				if (sum != 0)
				{
					double scale = scalars[i] / sum;
					for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i] + cnt_loc; ++j)
						Q_LocLoc_values[j] *= scale;
					for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i] + cnt_ext; ++j)
						Q_LocExt_values[j] *= scale;
				}
				else
				{
					for (int j = Q_LocLoc_rowptr[i]; j < Q_LocLoc_rowptr[i] + cnt_loc; ++j)
						Q_LocLoc_values[j] = 0.0;
					for (int j = Q_LocExt_rowptr[i]; j < Q_LocExt_rowptr[i] + cnt_ext; ++j)
						Q_LocExt_values[j] = 0.0;
				}

				P_LocLoc_rowptr[i + 1] = cnt_loc;
				P_LocExt_rowptr[i + 1] = cnt_ext;
			}
			else
			{
				P_LocLoc_rowptr[i + 1] = 1;
				P_LocExt_rowptr[i + 1] = 0;
			}
		}
	}

	for (int i = 0; i < n; ++i)
		P_LocLoc_rowptr[i + 1] += P_LocLoc_rowptr[i];
	for (int i = 0; i < n; ++i)
		P_LocExt_rowptr[i + 1] += P_LocExt_rowptr[i];

	int* P_LocLoc_colind = new int[P_LocLoc_rowptr[n]];
	double* P_LocLoc_values = new double[P_LocLoc_rowptr[n]];
	int* P_LocExt_colind = new int[P_LocExt_rowptr[n]];
	double* P_LocExt_values = new double[P_LocExt_rowptr[n]];

#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
	{
		for (int j = Q_LocLoc_rowptr[i], k = P_LocLoc_rowptr[i]; k < P_LocLoc_rowptr[i + 1]; ++j, ++k)
		{
			P_LocLoc_colind[k] = Q_LocLoc_colind[j];
			P_LocLoc_values[k] = Q_LocLoc_values[j];
		}

		for (int j = Q_LocExt_rowptr[i], k = P_LocExt_rowptr[i]; k < P_LocExt_rowptr[i + 1]; ++j, ++k)
		{
			P_LocExt_colind[k] = Q_LocExt_colind[j];
			P_LocExt_values[k] = Q_LocExt_values[j];
		}
	}

	int* P_neighbors = new int[Q_num_neighbors];
	for (int r = 0; r < Q_num_neighbors; ++r)
		P_neighbors[r] = Q_neighbors[r];

	int* w = new int[Q.LocExt.size[1]];
	for (int i = 0; i < Q.LocExt.size[1]; ++i)
		w[i] = -1;

	for (int j = P_LocExt_rowptr[0]; j < P_LocExt_rowptr[n]; ++j)
		w[P_LocExt_colind[j]] = 1;

	int* P_Ext_map = new int[Q.LocExt.size[1]];
	int* P_Ext_ptr = new int[Q_num_neighbors + 1];

	P_Ext_ptr[0] = 0;

	for (int r = 0, k = 0; r < Q_num_neighbors; P_Ext_ptr[++r] = k)
	{
		for (int i = Q_Ext_ptr[r]; i < Q_Ext_ptr[r + 1]; ++i)
		{
			if (w[i] == 1)
			{
				w[i] = k;
				P_Ext_map[k++] = Q_Ext_map[i];
			}
		}
	}

	for (int j = P_LocExt_rowptr[0]; j < P_LocExt_rowptr[n]; ++j)
		P_LocExt_colind[j] = w[P_LocExt_colind[j]];

	delete[] w;
	delete[] Ext_CF_map;
	delete[] marks;
	delete[] scalars;

	P.Free();

	P.comm = comm;

	P.LocLoc.size[0] = n;
	P.LocLoc.size[1] = num_coarse;
	P.LocLoc.rowptr = P_LocLoc_rowptr;
	P.LocLoc.colind = P_LocLoc_colind;
	P.LocLoc.values = P_LocLoc_values;

	P.LocExt.size[0] = n;
	P.LocExt.size[1] = P_Ext_ptr[Q_num_neighbors];
	P.LocExt.rowptr = P_LocExt_rowptr;
	P.LocExt.colind = P_LocExt_colind;
	P.LocExt.values = P_LocExt_values;

	P.num_neighbors = Q_num_neighbors;
	P.neighbor_ranks = P_neighbors;
	P.Ext_ptr = P_Ext_ptr;
	P.Ext_map = P_Ext_map;
}

}